from PIL import Image
import numpy as np
from torch.utils.data import Dataset
from torchvision import transforms
from os import path
import json

class DeepFashionDataset(Dataset):
    def __init__(self, data_path, train=True):
        self.data_path = data_path
        self.train = train
        self.img_dir = path.join(data_path, "images")
        self.transform = transforms.Compose(
    [
        transforms.Resize((224,224)),
        transforms.ToTensor(),
        transforms.Normalize([0.831, 0.813, 0.806],[0.229, 0.244, 0.252])
    ]
)
        if train:
            json_dir = path.join(self.data_path,'test_captions.json')
        else:
            json_dir = path.join(self.data_path,'train_captions.json')

        with open(json_dir, "r") as f:
            data = json.load(f)
        
        self.images_name = list(data.keys())
        self.labels = np.array(list(data.values()))

    def __len__(self):
        return len(self.images_name)

    def __getitem__(self, idx):
        img_name = self.images_name[idx]
        img = Image.open(path.join(self.img_dir, img_name))
        img = self.transform(img)
        return img, self.labels[idx]